# -*- coding: utf-8 -*-
"""
@Time    : 2025/3/4 20:45 
@Author  : ZhangShenao 
@File    : agent.py 
@Desc    : Agent模块
"""
import json
import re

from llm import LLM
from prompt import build_react_prompt
from tools import TOOLS_DESCRIPTION, query_fruit_unit_price

if __name__ == '__main__':
    # 构造Prompt
    instructions = "你是一个居家生活小助手，可以回答用户的日常问题。"
    query = "我想买2个苹果和3根香蕉，一共需要多少钱？"
    prompt = build_react_prompt(instructions=instructions,
                                query=query,
                                tool_desc=TOOLS_DESCRIPTION,
                                tool_name="query_fruit_unit_price",
                                )

    # 创建LLM
    llm = LLM(model_name="deepseek-v3")

    # 保存上下文
    print(f"初始提问: {prompt}")
    messages = [{"role": "user", "content": prompt}]

    # 执行ReAct过程
    while True:
        response = llm.send_msg(messages)
        response_text = response.choices[0].message.content

        print(f"大模型的回复：\n{response_text}")

        # 通过正则表达式匹配,判断是否结束执行
        final_answer_match = re.search(r'Final Answer:\s*(.*)', response_text)
        if final_answer_match:
            final_answer = final_answer_match.group(1)
            print("最终答案:", final_answer)
            break

        # 保存上下文
        messages.append(response.choices[0].message)

        # 通过正则表达式匹配,解析Function Calling参数
        action_match = re.search(r'Action:\s*(\w+)', response_text)
        action_input_match = re.search(r'Action Input:\s*({.*?}|".*?")', response_text, re.DOTALL)

        # 匹配需要调用的工具
        if action_match and action_input_match:
            tool_name = action_match.group(1)
            params = json.loads(action_input_match.group(1))
            print(f"需要执行Function Calling, 工具名称: {tool_name}, 调用参数: {params}")

            # 调用工具,获取执行结果
            if tool_name == "query_fruit_unit_price":
                observation = query_fruit_unit_price(params['fruit_name'])
                print(f"工具的执行结果: \n{observation}", )

                # 保存上下文
                messages.append({"role": "user", "content": f"Observation: {observation}"})
